[WIP] perf(policies): Make ACT policy compatible with torch.compile
#2159
+751
−2
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Adds `benchmark_inference_compile_lerobot.py` from https://gist.github.com/AdilZouitine/3574664e4cf71605986b49e9148d29ab.
…pile graph break Removed .item() calls from loss_dict in forward() to avoid breaking the torch.compile computation graph. The tensor-to-scalar conversion is now handled in the training script instead.
In the inference benchmark in `benchmark_inference_compile_lerobot.py`, the `test_correctness` method failed to properly compare compiled `ACTPolicy` inference results. This was due to `policy_original` and `policy_compiled` sharing the same `_action_queue` object. Previously, the call order was: 1. `policy_original.reset()` 2. `policy_compiled.reset()` 3. `policy_original.select_action()` 4. `policy_compiled.select_action()` Because the `_action_queue` is shared, `policy_original.select_action()` would run inference (`predict_action_chunk`) and extend the queue with `n_action_steps` actions. `policy_compiled.select_action()` would then find a non-empty queue and simply pop an action, bypassing its own compiled inference logic. This commit reorders the calls to: 1. `policy_original.reset()` 2. `policy_original.select_action()` 3. `policy_compiled.reset()` 4. `policy_compiled.select_action()` This change ensures that `policy_compiled.reset()` clears the shared queue *after* the original policy's action selection. Consequently, `policy_compiled.select_action()` finds an empty queue and executes its own compiled inference, allowing for a correct comparison. With this fix, the compiled `ACTPolicy` inference check within `test_correctness` now passes, validating that the compiled inference output matches the original.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Note
Work in Progress !
This PR is my attempt to address issue #2061. As it's a complex issue and the PR is still under development, all collaboration, help, and suggestions would be greatly appreciated!
What this does
(WIP) Explain what this PR does.
src/lerobot/policies/act/modeling_act.pysrc/lerobot/scripts/lerobot_train.pywith accelerator.autocast(): loss, output_dict = policy.forward(batch) + output_dict = {k: v.item() if isinstance(v, torch.Tensor) else v for k, v in output_dict.items()}Current progress
✅ Compilation: PASSED
❌ Correctness: FAILED
❌ Performance: FAILED
✅ Benchmarking: PASSED
When compiling the policy with
torch.compile(policy, mode="default"), theloss_diffof training test withintest_correctnessfails to meet the1e-5tolerance. Using eithertorch.compile(policy, mode="default", options={"fallback_random": True})ortorch.compile(policy, mode="default", backend="aot_eager")satisfies the tolerance, but neither configuration provides a noticeable performance improvement.How it was tested
(WIP) Explain/show how you tested your changes.
How to checkout & try? (for the reviewer)
(WIP) Provide a simple way for the reviewer to try out your changes.